import numpy as np
import imageio
from image_utils import logging


def split_image_channels(image):
    '''
    该函数用于分割图像的颜色通道
    :param image: 输入图像，numpy数组
    :return: 图像的颜色通道列表
    '''
    try:
        logging.info("开始分割图像颜色通道")

        # 1. 检测多通道图像 image 是否真多通道的，如果是，显示错误信息并返回
        if image.ndim < 3 or image.shape[-1] < 2:
            raise ValueError("图像通道数必须大于等于2")

        # 2. 使用numpy的split函数，沿着通道维度分割图像
        channel_list = np.split(image, image.shape[-1], axis=-1)

        # 3. 把分割后的通道图像添加到一个列表 channel_list 中，注意要去掉多余的维度
        channel_list = [channel.squeeze() for channel in channel_list]

        logging.info("成功分割图像颜色通道")
        return channel_list
    except Exception as e:
        raise ValueError(f"在分割图像颜色通道时发生错误: {e}")


def merge_image_channels(file_paths):
    '''
    该函数用于合并多个单通道图像，创建一个多通道图像
    :param file_paths: 包含单通道图像文件路径的列表
    :return: 合并后的多通道图像
    '''
    try:
        logging.info("开始合并图像通道")

        # 初始化一个空的numpy数组用于存储合并后的图像
        merged_image = None

        # 遍历每个文件路径
        for file_path in file_paths:
            # 加载图像
            channel_image = imageio.imread(file_path)

            # 检查每个图像是否是单通道的
            if channel_image.ndim != 2:
                raise ValueError("要合并的子图像必须是单通道图像，不能是多通道图像")

            # 如果是第一个通道，初始化merged_image
            if merged_image is None:
                merged_image = np.zeros((channel_image.shape[0], channel_image.shape[1], len(file_paths)),
                                        dtype=channel_image.dtype)

            # 将当前通道图像添加到merged_image中
            merged_image[..., file_paths.index(file_path)] = channel_image

        logging.info("成功合并图像通道")
        return merged_image
    except Exception as e:
        raise ValueError(f"在合并图像通道时发生错误: {e}")